#include "RTPUnpacker.h"

#include "../inc/Utils.h"
#include "VideoDecoder.h"

#define RTP_VERSION 2  

RTPUnpacker::RTPUnpacker(unsigned char H264PAYLOADTYPE = 96) :
		m_bSPSFound(false), m_bWaitKeyFrame(true), m_bPrevFrameEnd(false), m_bAssemblingFrame(
				false), m_wSeq(1234), m_ssrc(0)
{
	m_pBuf = new char[MAX_FRAME_SIZE];
	if (m_pBuf == NULL)
	{
		return;
	}

	m_H264PAYLOADTYPE = H264PAYLOADTYPE;
	m_pEnd = m_pBuf + MAX_FRAME_SIZE;
	m_pStart = m_pBuf;
	m_dwSize = 0;

	receivedPacket = 0;
	lostPacket = 0;
}

RTPUnpacker::~RTPUnpacker(void)
{
	delete[] m_pBuf;
}

// nSize:size of the rtp packet. video data may be corrupted
char* RTPUnpacker::Parse_RTP_Packet(char *pBuf, unsigned short nSize,
		int *outSize)
{
	if (nSize <= 12)
	{
		return NULL;
	}

	receivedPacket++;
	char *cp = (char*) &m_RTP_Header;
	cp[0] = pBuf[0];
	cp[1] = pBuf[1];

	m_RTP_Header.seq = pBuf[2];
	m_RTP_Header.seq <<= 8;
	m_RTP_Header.seq |= pBuf[3];

	m_RTP_Header.ts = pBuf[4];
	m_RTP_Header.ts <<= 8;
	m_RTP_Header.ts |= pBuf[5];
	m_RTP_Header.ts <<= 8;
	m_RTP_Header.ts |= pBuf[6];
	m_RTP_Header.ts <<= 8;
	m_RTP_Header.ts |= pBuf[7];

	m_RTP_Header.ssrc = pBuf[8];
	m_RTP_Header.ssrc <<= 8;
	m_RTP_Header.ssrc |= pBuf[9];
	m_RTP_Header.ssrc <<= 8;
	m_RTP_Header.ssrc |= pBuf[10];
	m_RTP_Header.ssrc <<= 8;
	m_RTP_Header.ssrc |= pBuf[11];

	char *pPayload = pBuf + 12;
	int PayloadSize = nSize - 12;

	// Check the RTP version number (it should be 2):
	if (m_RTP_Header.v != RTP_VERSION)
	{
		return NULL;
	}

	// Check the Payload Type.
	if (m_RTP_Header.pt < m_H264PAYLOADTYPE)
	{
		return NULL;
	}

	int PayloadType = pPayload[0] & 0x1f;
	int NALType = PayloadType;
	if (NALType == 28) // FU_A
	{
		if (PayloadSize < 2)
		{
			return NULL;
		}

		NALType = pPayload[1] & 0x1f;
	}

	if (m_ssrc != m_RTP_Header.ssrc)
	{
		m_ssrc = m_RTP_Header.ssrc;
		SetLostPacket();
	}

	if (NALType == 0x07) // SPS
	{
		m_bSPSFound = true;
	}

	if (!m_bSPSFound)
	{
		return NULL;
	}

	if (NALType == 0x07 || NALType == 0x08) // SPS PPS
	{
		m_wSeq = m_RTP_Header.seq;
		m_bPrevFrameEnd = true;

		pPayload -= 4;
		*((char*) (pPayload)) = 0x01000000;
		*outSize = PayloadSize + 4;
		return pPayload;
	}

	//if ( m_bWaitKeyFrame )
	//{
	//    if ( m_RTP_Header.m ) // frame end
	//    {
	//        m_bPrevFrameEnd = true ;
	//        if ( !m_bAssemblingFrame )
	//        {
	//            m_wSeq = m_RTP_Header.seq ;
	//            return NULL ;
	//        }
	//    }

	//    if ( !m_bPrevFrameEnd )
	//    {
	//        m_wSeq = m_RTP_Header.seq ;
	//        return NULL ;
	//    }
	//    else
	//    {
	//        if ( NALType != 0x05 ) // KEY FRAME
	//        {
	//            m_wSeq = m_RTP_Header.seq ;
	//            m_bPrevFrameEnd = false ;
	//            return NULL ;
	//        }
	//    }
	//}

///////////////////////////////////////////////////////////////  

	if (m_RTP_Header.seq != (int) (m_wSeq + 1)) // lost packet
	{
		if (m_RTP_Header.seq - m_wSeq > 0)
			lostPacket += (m_RTP_Header.seq - m_wSeq);
		else
			lostPacket += (m_RTP_Header.seq - m_wSeq + 65536);
		_print_log("received packet:%d, packet lost:%d, loss rate:%d%%\n",
				receivedPacket, lostPacket,
				lostPacket * 100 / (receivedPacket + lostPacket));
		m_wSeq = m_RTP_Header.seq;
		//SetLostPacket () ;
		//return NULL ;
		if (m_RTP_Header.seq - m_wSeq > 5)
			return NULL;
	}
	else
	{
		// no loss
	}
	m_wSeq = m_RTP_Header.seq;
	m_bAssemblingFrame = true;

	if (PayloadType != 28) // whole NAL
	{
		*((int*) (m_pStart)) = 0x01000000;
		m_pStart += 4;
		m_dwSize += 4;
	}
	else // FU_A
	{
		if (pPayload[1] & 0x80) // FU_A start
		{
			*((int*) (m_pStart)) = 0x01000000;
			m_pStart += 4;
			m_dwSize += 4;

			pPayload[1] = (pPayload[0] & 0xE0) | NALType;

			pPayload += 1;
			PayloadSize -= 1;
		}
		else
		{
			pPayload += 2;
			PayloadSize -= 2;
		}
	}

	if (m_pStart + PayloadSize < m_pEnd)
	{
		memcpy(m_pStart, pPayload, PayloadSize);
		m_dwSize += PayloadSize;
		m_pStart += PayloadSize;
	}
	else // memory overflow
	{
		_print_log("overflow\n");
		SetLostPacket();
		return NULL;
	}

	if (m_RTP_Header.m) // frame end
	{
		*outSize = m_dwSize;

		m_pStart = m_pBuf;
		m_dwSize = 0;

		if (NALType == 0x05) // KEY FRAME
		{
			m_bWaitKeyFrame = false;
		}
		return m_pBuf;
	}
	else
	{
		return NULL;
	}

}

void RTPUnpacker::SetLostPacket()
{
	m_bSPSFound = false;
	m_bWaitKeyFrame = true;
	m_bPrevFrameEnd = false;
	m_bAssemblingFrame = false;
	m_pStart = m_pBuf;
	m_dwSize = 0;
}

